package com.unisound.tracelog.client.core.service.impl;

import cn.hutool.core.lang.generator.SnowflakeGenerator;
import com.unisound.tracelog.client.core.context.SpanContext;
import com.unisound.tracelog.client.core.context.SpanContextManager;
import com.unisound.tracelog.client.core.listener.SpanEvent;
import com.unisound.tracelog.client.core.listener.TagEvent;
import com.unisound.tracelog.client.core.listener.TraceLogListeners;
import com.unisound.tracelog.client.core.model.RootStartSpan;
import com.unisound.tracelog.client.core.model.StartSpan;
import com.unisound.tracelog.client.core.properties.TraceLogClientProperties;
import com.unisound.tracelog.client.core.service.TraceLogService;
import com.unisound.tracelog.client.core.support.parameter.ParameterValue;
import com.unisound.tracelog.common.model.span.EndSpan;
import com.unisound.tracelog.common.model.span.Span;
import com.unisound.tracelog.common.model.span.StartSpanResult;
import com.unisound.tracelog.common.model.span.Tag;
import com.unisound.tracelog.common.utils.CommonUtils;
import com.unisound.tracelog.common.utils.JsonUtils;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.*;

public class TraceLogServiceImpl implements TraceLogService {

    private TraceLogClientProperties traceLogClientProperties;

    private SnowflakeGenerator traceLogSnowflakeGenerator;

    private TraceLogListeners traceLogListeners;

    public TraceLogServiceImpl(TraceLogClientProperties traceLogClientProperties, SnowflakeGenerator traceLogSnowflakeGenerator, TraceLogListeners traceLogListeners) {
        this.traceLogClientProperties = traceLogClientProperties;
        this.traceLogSnowflakeGenerator = traceLogSnowflakeGenerator;
        this.traceLogListeners = traceLogListeners;
    }

    @Override
    public StartSpanResult start(StartSpan startSpan) {
        if (startSpan == null ||
            StringUtils.isEmpty(startSpan.getName())) return null;

        long startTime = System.currentTimeMillis();

        String name = startSpan.getName();

        String serverName = startSpan.getServerName();
        if (StringUtils.isEmpty(serverName)) serverName = traceLogClientProperties.getDefaultServerName();
        if (StringUtils.isNotEmpty(serverName)) serverName = serverName.toUpperCase(Locale.ROOT);

        Integer type = startSpan.getType();

        Long traceId = 0L;
        Long parentSpanId = 0L;
        Long spanId = 0L;

        SpanContext spanContext = SpanContextManager.getSpanContext();
        boolean root = SpanContextManager.isRoot();
        if (root) {
            traceId = traceLogSnowflakeGenerator.next();
            spanId = traceId;
        } else {
            traceId = spanContext.getTraceId();
            spanId = traceLogSnowflakeGenerator.next();
            parentSpanId = spanContext.getSpanId();
        }

        String request = toJson(startSpan.getParameterMap());

        traceLogListeners.publishSpanEvent(new SpanEvent(Span.startSpan(spanId, traceId, parentSpanId, name, serverName, type, request)));

        if (root) {
            if (startSpan instanceof RootStartSpan) {
                RootStartSpan rootStartSpan = (RootStartSpan) startSpan;
                List<Tag> tags = tagList(rootStartSpan.getTagMap(), traceId);
                if (CollectionUtils.isNotEmpty(tags)) {
                    traceLogListeners.publishTagEvent(new TagEvent(tags));
                }
            }
        }

        SpanContextManager.add(new SpanContext(traceId, spanId, startSpan.getName(), parentSpanId));

        return new StartSpanResult(traceId, parentSpanId, spanId, name, serverName, type, request, startTime);
    }

    private List<Tag> tagList(Map<String, String> tagMap, Long traceId) {
        if (tagMap == null || tagMap.isEmpty()) return null;
        if (traceId == null || traceId <= 0) return null;

        List<Tag> tags = new ArrayList<>();
        for (Map.Entry<String, String> entry : tagMap.entrySet()) {
            String tagKey = entry.getKey();
            String tagValue = entry.getValue();
            if (StringUtils.isEmpty(tagKey) || StringUtils.isEmpty(tagValue)) continue;

            tags.add(new Tag(traceLogSnowflakeGenerator.next(), traceId, tagKey, tagValue));
        }

        return tags;
    }

    @Override
    public void end(EndSpan endSpan) {
        if (endSpan == null ||
            endSpan.getStartSpanResult() == null ||
            endSpan.getStartSpanResult().getSpanId() == null ||
            endSpan.getStartSpanResult().getTraceId() == null ||
            endSpan.getStartSpanResult().getParentSpanId() == null ||
            StringUtils.isEmpty(endSpan.getStartSpanResult().getName())) return;

        StartSpanResult startSpanResult = endSpan.getStartSpanResult();
        Long spanId = startSpanResult.getSpanId();
        Long traceId = startSpanResult.getTraceId();
        Long parentSpanId = startSpanResult.getParentSpanId();
        String name = startSpanResult.getName();
        String serverName = startSpanResult.getServerName();
        Integer type = startSpanResult.getType();
        Object response = endSpan.getResult();
        Map<String, String> errorMessageMap = endSpan.getErrorMessageMap();
        String jsonResult = CommonUtils.jsonResult(response, errorMessageMap);

        Span currEndSpan = Span.endSpan(spanId, traceId, parentSpanId, name, serverName, type, jsonResult);
        currEndSpan.setResultFlag(endSpan.getResultFlag());
        currEndSpan.setElapsedTimeMillis(currEndSpan.getExecTime() - startSpanResult.getStartTime());
        traceLogListeners.publishSpanEvent(new SpanEvent(currEndSpan));

        SpanContextManager.remove();
    }

    @Override
    public boolean support() {
        return traceLogClientProperties.isEnable();
    }

    private String toJson(Map<String, ParameterValue> parameterMap) {
        if (parameterMap == null || parameterMap.isEmpty()) return null;

        Map<String, Object> map = new HashMap<>();
        parameterMap.forEach((key, value) -> {
            Object object = value.getObject();
            if (object != null) {
                if (value.isUseToString()) {
                    map.put(key, object.toString());
                } else {
                    map.put(key, object);
                }
            }
        });

        if (map.isEmpty()) return null;
        return JsonUtils.toJson(map);
    }

}
